# USAGE
# python scan.py --image .\test.jpg --output output/scanned --format jpg
'''
读取手机拍摄的照片，将其处理成扫描后的图片，同时保留红色印章颜色
如果处理后图片有部分噪点，可以先将图片对比度调高，再进行程序处理
如果待扫描的照片占据整个屏幕，没有明显的分隔界面，可以先手动裁剪图片或添加不小于2像素粗细的边框
微调参数在configure_red_seal_detection()函数中

Reads a photo taken by a mobile phone, processes it into a scanned image, and preserves the color of the red seal.
If the processed image has some noise, you can first adjust the contrast of the image and then process it with the program.
If the photo to be scanned occupies the entire screen and has no clear separation interface, you can first manually crop the image or add a border of at least 2 pixels in thickness.
Fine-tune parameters in the configure_red_seal_detection() function.

This project is based on https://pyimagesearch.com/2014/09/01/build-kick-ass-mobile-document-scanner-just-5-minutes/
'''

from pyimagesearch.transform import four_point_transform
from skimage.filters import threshold_local
import numpy as np
import argparse
import cv2
import imutils
import os

ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True, help="Path to the image to be scanned")
ap.add_argument(
    "-o",
    "--output",
    default="scanned",
    help="Path to output scanned image (without extension)",
)
ap.add_argument(
    "-f",
    "--format",
    default="jpg",
    choices=["jpg", "png", "pdf"],
    help="Output format: jpg, png or pdf",
)
args = vars(ap.parse_args())

# 加载图像并计算旧高度与新高度的比率，复制图像并调整大小
image = cv2.imread(args["image"])
ratio = image.shape[0] / 500.0
orig = image.copy()
image = imutils.resize(image, height=500)

# 将图像转换为灰度图，进行模糊处理，并在图像中找到边缘
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gray = cv2.GaussianBlur(gray, (5, 5), 0)
edged = cv2.Canny(gray, 75, 200)

# 显示原始图像和边缘检测后的图像
print("STEP 1: Edge Detection")
cv2.imshow("Image", image)
cv2.imshow("Edged", edged)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 在边缘图像中找到轮廓，只保留最大的轮廓，并初始化屏幕轮廓
cnts = cv2.findContours(edged.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
cnts = sorted(cnts, key=cv2.contourArea, reverse=True)[:5]

# 遍历轮廓
for c in cnts:
    # 近似轮廓
    peri = cv2.arcLength(c, True)
    approx = cv2.approxPolyDP(c, 0.02 * peri, True)

    # 如果我们的近似轮廓有四个点，那么我们可以假设已经找到了屏幕
    if len(approx) == 4:
        screenCnt = approx
        break

# 显示纸张的轮廓（轮廓）
print("STEP 2: Find contours of paper")
cv2.drawContours(image, [screenCnt], -1, (0, 255, 0), 2)
cv2.imshow("Outline", image)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 应用四点变换以获得原始图像的俯视图
warped = four_point_transform(orig, screenCnt.reshape(4, 2) * ratio)

# 保留原始变形图像的副本以供颜色提取
warped_color = warped.copy()

# 将变形图像转换为灰度图以进行文档处理
warped_gray = cv2.cvtColor(warped, cv2.COLOR_BGR2GRAY)
T = threshold_local(warped_gray, 11, offset=10, method="gaussian")
warped_bw = (warped_gray > T).astype("uint8") * 255


def configure_red_seal_detection(
    # HSV范围参数
    h_low1=0,  # 第一个红色范围的色相下限 (0-180)
    h_high1=25,  # 第一个红色范围的色相上限 (0-180)
    h_low2=150,  # 第二个红色范围的色相下限 (0-180)
    h_high2=180,  # 第二个红色范围的色相上限 (0-180)
    s_low=25,  # 饱和度下限 (0-255) - 值越低捕获的红色越浅
    s_high=255,  # 饱和度上限 (0-255)
    v_low=25,  # 亮度下限 (0-255) - 值越低捕获的红色越暗
    v_high=255,  # 亮度上限 (0-255)
    # 降噪参数
    blur_kernel=(3, 3),  # 高斯模糊核大小 - 值越大，去噪效果越强但可能丢失细节
    blur_sigma=0,  # 高斯模糊的标准差 - 值越大，模糊效果越强
    # 形态学操作参数
    morph_kernel=(3, 3),  # 形态学操作核大小 - 值越大，对印章形状的修改越明显
    debug=False,  # 是否显示红色掩码的调试窗口
):
    """
    配置印章检测的参数

    Args:
    ----------
    HSV范围参数:
        h_*: 色相范围 (0-180)
            - 红色在HSV色彩空间中横跨0度和180度附近
            - 因此需要两组范围: (h_low1, h_high1) 和 (h_low2, h_high2)
            - 增大范围可以捕获更多的红色色调

        s_*: 饱和度范围 (0-255)
            - 饱和度表示颜色的纯度
            - 降低s_low可以捕获更浅的红色

        v_*: 亮度范围 (0-255)
            - 控制颜色的明暗
            - 降低v_low可以捕获更暗的红色

    降噪参数:
        blur_kernel: 高斯模糊核大小
            - 必须是奇数
            - 值越大，去噪效果越强，但会丢失细节

        blur_sigma: 高斯模糊标准差
            - 0表示自动计算
            - 值越大，模糊效果越强

    形态学参数:
        morph_kernel: 形态学操作核大小
            - 值越大，对印章形状的修改越明显
            - 可以帮助填充印章中的小空隙或去除噪点

    Returns:
    -------
    dict: 包含所有配置参数的字典
    """

    config = {
        'lower_red1': np.array([h_low1, s_low, v_low]),
        'upper_red1': np.array([h_high1, s_high, v_high]),
        'lower_red2': np.array([h_low2, s_low, v_low]),
        'upper_red2': np.array([h_high2, s_high, v_high]),
        'blur_kernel': blur_kernel,
        'blur_sigma': blur_sigma,
        'morph_kernel': morph_kernel,
        'debug': debug,
    }
    return config


def process_red_seal(warped_color, config):
    """
    使用给定的配置处理图像中的红色印章

    Args:
    ------
    warped_color: ndarray
        原始彩色图像
    config: dict
        由configure_red_seal_detection()返回的配置字典

    Returns:
    --------
    tuple: (red_mask, red_regions)
        red_mask: 红色区域的二值掩码
        red_regions: 提取出的红色区域
    """
    # 转换为HSV颜色空间
    warped_hsv = cv2.cvtColor(warped_color, cv2.COLOR_BGR2HSV)

    # 创建红色区域的掩码
    mask1 = cv2.inRange(warped_hsv, config['lower_red1'], config['upper_red1'])
    mask2 = cv2.inRange(warped_hsv, config['lower_red2'], config['upper_red2'])
    red_mask = cv2.bitwise_or(mask1, mask2)

    # 应用高斯模糊
    red_mask = cv2.GaussianBlur(red_mask, config['blur_kernel'], config['blur_sigma'])

    # 使用形态学操作优化掩码
    kernel = np.ones(config['morph_kernel'], np.uint8)
    red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel)
    red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel)

    # 如果启用调试，显示调试窗口
    if config['debug']:
        cv2.imshow("Red Mask", imutils.resize(red_mask, height=650))

    # 提取红色区域
    red_regions = cv2.bitwise_and(warped_color, warped_color, mask=red_mask)

    return red_mask, red_regions


# 配置印章检测参数
seal_config = configure_red_seal_detection(
    # 调整参数
    h_low1=0,
    h_high1=20,
    h_low2=140,
    h_high2=180,
    s_low=25,
    v_low=10,
    # 如果印章边缘有噪点，可以增加模糊核的大小，必须是奇数
    blur_kernel=(1, 1),
    # 如果印章有断开，可以增加形态学核的大小
    morph_kernel=(0, 0),
    debug=True,
)

# 处理红色印章
red_mask, red_regions = process_red_seal(warped_color, seal_config)

# 将warped_bw转换为3通道以与颜色组合
warped_bw_3ch = cv2.cvtColor(warped_bw, cv2.COLOR_GRAY2BGR)

# 将黑白文档与红色区域组合
warped_final = cv2.bitwise_and(
    warped_bw_3ch, warped_bw_3ch, mask=cv2.bitwise_not(red_mask)
)
warped_final = cv2.add(warped_final, red_regions)

# 显示原始图像和扫描后的图像
print("STEP 3: Apply perspective transform")
cv2.imshow("Original", imutils.resize(orig, height=650))
cv2.imshow("Scanned", imutils.resize(warped_final, height=650))

# 保存扫描后的图像
output_path = args["output"]
output_format = args["format"].lower()
if not os.path.exists(os.path.dirname(output_path)):
    os.makedirs(os.path.dirname(output_path))

if output_format in ["jpg", "png"]:
    # 保存为图像
    output_file = f"{output_path}.{output_format}"
    cv2.imwrite(output_file, warped_final)
    print(f"Saved scanned image as: {output_file}")
elif output_format == "pdf":
    # 保存为PDF
    from PIL import Image

    # 将OpenCV图像转换为PIL图像（BGR到RGB）
    warped_final_rgb = cv2.cvtColor(warped_final, cv2.COLOR_BGR2RGB)
    pil_image = Image.fromarray(warped_final_rgb)

    # 保存为PDF
    output_file = f"{output_path}.pdf"
    pil_image.save(output_file, "PDF", resolution=100.0)
    print(f"Saved scanned image as: {output_file}")

cv2.waitKey(0)  # press any key to exit
cv2.destroyAllWindows()
